#include "jpeg_globals.h"
#include "jpeg_types.h"
#include "jpeg_markers.h"
#include "jpeg_dct.h"
#include "jpeg_quantize.h"
#include "jpeg_vlc.h"
#include "jpeg_bitstream.h"

void memcpy(SBYTE *dst, BYTE *src, int size)
{
	for (; size; size--)
		*dst++=*src++;
}

typedef struct image_t_
{
	BYTE *p_buf_y;
	BYTE *p_buf_u;
	BYTE *p_buf_v;
	DWORD yuv_format;
}image_t;

BYTE *g_stream_buf_start;

bitstring YDC_HT[12];
bitstring CbDC_HT[12];
bitstring YAC_HT[256];
bitstring CbAC_HT[256];
BYTE *category_alloc;
BYTE *category;
bitstring *bitcode_alloc;
bitstring *bitcode; // their bitcoded representation

SBYTE YDU[64*4]; // This is the Data Unit of Y after YCbCr->RGB transformation
SBYTE CbDU[64];
SBYTE CrDU[64];
SWORD DU_DCT[64]; // Current DU (after DCT and quantization) which we'll zigzag
SWORD DU[64]; //zigzag reordered DU which will be Huffman coded

BYTE restart_interval,restart_cnt,restart_mod;
SBYTE restart_en,restart_flag,restart_ini;

colorRGB *RGB_buffer; //image to be encoded
DWORD g_lum_stride, g_cb_stride;
SDWORD g_mcuNum; //there is how many MCU in the image
SDWORD g_mcuYcnt;//there is how many y blocks in a mcu
DWORD jpgEncFormat;
image_t g_image;
BYTE quantity=32;

static SDWORD YRtab[256],YGtab[256],YBtab[256];
static SDWORD CbRtab[256],CbGtab[256],CbBtab[256];
static SDWORD CrRtab[256],CrGtab[256],CrBtab[256];

DWORD mcu_num_x, mcu_num_y;
BYTE mcu_size_x, mcu_size_y;

void init_jpeg_encoder()
{
	jpg_init_bitstream();
	jpg_set_DQTinfo(quantity);
	jpg_set_DHTinfo();
	init_Huffman_tables(YDC_HT, CbDC_HT, YAC_HT, CbAC_HT);
}

void load_data_units_for_yuv (DWORD mcu_count)
{
	BYTE *ptrY, *ptrU, *ptrV;
	int i, j;
	DWORD xpos, ypos;
	if(jpgEncFormat == 0) //444
	{
		xpos = (mcu_count%(g_lum_stride/8))*8;
		ypos = (mcu_count/(g_lum_stride/8))*8;
		ptrY = g_image.p_buf_y + xpos + ypos*g_lum_stride;
		for(i=0; i<8; i++)
		{
			memcpy(&YDU[i*8], ptrY, 8);
			ptrY += g_lum_stride;
		}
	}
	else if(jpgEncFormat == 1) //422
	{
		xpos = (mcu_count%(g_lum_stride/16))*16;
		ypos = (mcu_count/(g_lum_stride/16))*8;
		for(j=0;j<2;j++)
		{
			ptrY = g_image.p_buf_y + xpos + ypos*g_lum_stride + j*8;
			for(i=0; i< 8; i++)
			{
				memcpy(&YDU[i*8 + j*64], ptrY, 8);
				ptrY += g_lum_stride;
			}
		}
	}
	else if(jpgEncFormat == 2) //411
	{
		xpos = (mcu_count%(g_lum_stride/32))*32;
		ypos = (mcu_count/(g_lum_stride/32))*8;
		for(j=0;j<4;j++)
		{
			ptrY = g_image.p_buf_y + xpos + ypos*g_lum_stride + j*8;
			for(i=0; i< 8; i++)
			{
				memcpy(&YDU[i*8 + j*64], ptrY, 8);
				ptrY += g_lum_stride;
			}
		}
	}
	else //420
	{
		xpos = (mcu_count%mcu_num_x)*mcu_size_x;
		ypos = (mcu_count/mcu_num_x)*mcu_size_y;
		for(j=0;j<4;j++)
		{
			ptrY = g_image.p_buf_y + xpos + (ypos+(j/2)*8)*g_lum_stride + (j%2)*8;
			for(i=0; i< 8; i++)
			{
				memcpy(&YDU[i*8 + j*64], ptrY, 8);
				ptrY += g_lum_stride;
			}
		}
	}
	xpos = (mcu_count%(g_cb_stride/8))*8;
	ypos = (mcu_count/(g_cb_stride/8))*8;
	ptrU = g_image.p_buf_u + xpos + ypos*g_cb_stride;
	ptrV = g_image.p_buf_v + xpos + ypos*g_cb_stride;
	for(i=0; i< 8; i++)
	{
		memcpy(&CbDU[i*8], ptrU, 8);
		memcpy(&CrDU[i*8], ptrV, 8);
		ptrU += g_cb_stride;
		ptrV += g_cb_stride;
	}
}

void test_enc()
{
	int j, bkCnt;
	SWORD DCY=0,DCCb=0,DCCr=0;
	DWORD stream_len;
	WORD Ximage_original, Yimage_original;
	
	Ximage_original = 480;
	Yimage_original = 640;
	g_image.yuv_format = 3;
	jpgEncFormat = 3;
	mcu_size_x=16;
	mcu_size_y=16;
	g_lum_stride = Ximage_original;
	g_cb_stride = Ximage_original/2;
	
	g_image.p_buf_y = (BYTE*)0x5000000;
	g_image.p_buf_u = g_image.p_buf_y+Ximage_original*Yimage_original;
	g_image.p_buf_v = g_image.p_buf_u+Ximage_original*Yimage_original/4;
	g_stream_buf_start = (BYTE*)0x4000000;
	category_alloc = (BYTE*)0x3800000;
	bitcode_alloc = (bitstring*)0x3880000;
	mcu_num_x=Ximage_original/mcu_size_x;
	mcu_num_y=Yimage_original/mcu_size_y;
	
	g_mcuNum=mcu_num_x*mcu_num_y;
	g_mcuYcnt=4;

	init_jpeg_encoder();
	writeword(0xFFD8); //SOI
	jpg_write_APP0info();

	jpg_write_SOF0info(Ximage_original, Yimage_original, jpgEncFormat);

	jpg_write_DQTinfo();
	jpg_write_DHTinfo();

	jpg_write_SOSinfo();

	for (j=0; j<=g_mcuNum-1;j++)
	{
		load_data_units_for_yuv(j);

		//encode the Y data
		for(bkCnt=0;bkCnt<g_mcuYcnt;bkCnt++)
		{
			jpg_fdct(YDU+64*bkCnt, DU_DCT);
			jpg_quantize(DU_DCT, luminance_qt);
			jpg_vlc_data_unit(YDU,&DCY,YDC_HT,YAC_HT,0);
		}
	    //encode the U data
		jpg_fdct(CbDU, DU_DCT);
		jpg_quantize(DU_DCT, chrominance_qt);
		jpg_vlc_data_unit(CbDU,&DCCb,CbDC_HT,CbAC_HT,0);

		//encode the V data
		jpg_fdct(CrDU, DU_DCT);
		jpg_quantize(DU_DCT, chrominance_qt);
	    jpg_vlc_data_unit(CrDU,&DCCr,CbDC_HT,CbAC_HT,0);
	}
	
	jpg_byte_align_bitstream();
	jpg_write_EOIinfo();

	stream_len = jpg_get_streamlen();
}
